#!/usr/bin/env python3
"""
GPU 性能监控脚本 - 专为 RTX 3080 优化
"""
import torch
import time
import sys

def check_gpu_status():
    """检查 GPU 状态和性能"""
    print("🔍 GPU 环境检测")
    print("=" * 50)
    
    # 基础信息
    print(f"🐍 Python版本: {sys.version}")
    print(f"🔥 PyTorch版本: {torch.__version__}")
    
    # CUDA 支持
    cuda_available = torch.cuda.is_available()
    print(f"🚀 CUDA可用: {'✅ 是' if cuda_available else '❌ 否'}")
    
    if cuda_available:
        print(f"📌 CUDA版本: {torch.version.cuda}")
        print(f"🔧 cuDNN版本: {torch.backends.cudnn.version()}")
        
        # GPU 信息
        device_count = torch.cuda.device_count()
        print(f"🖥️  GPU数量: {device_count}")
        
        for i in range(device_count):
            gpu_name = torch.cuda.get_device_name(i)
            gpu_props = torch.cuda.get_device_properties(i)
            memory_total = gpu_props.total_memory / 1024**3
            
            print(f"\n📊 GPU {i}: {gpu_name}")
            print(f"   💾 总显存: {memory_total:.1f}GB")
            print(f"   🧮 计算能力: {gpu_props.major}.{gpu_props.minor}")
            print(f"   🔧 多处理器: {gpu_props.multi_processor_count}")
            
            # 当前显存使用
            torch.cuda.set_device(i)
            memory_allocated = torch.cuda.memory_allocated() / 1024**3
            memory_reserved = torch.cuda.memory_reserved() / 1024**3
            memory_free = memory_total - memory_reserved
            
            print(f"   📈 已分配: {memory_allocated:.2f}GB")
            print(f"   📦 已保留: {memory_reserved:.2f}GB") 
            print(f"   🆓 可用: {memory_free:.2f}GB")
            
        # 性能测试
        print(f"\n⚡ GPU 性能测试...")
        device = torch.device('cuda:0')
        
        # 创建测试张量
        size = 5000
        print(f"🧪 创建 {size}x{size} 矩阵进行测试...")
        
        start_time = time.time()
        a = torch.randn(size, size, device=device)
        b = torch.randn(size, size, device=device)
        
        # 矩阵乘法测试
        torch.cuda.synchronize()  # 等待GPU操作完成
        compute_start = time.time()
        
        c = torch.mm(a, b)
        torch.cuda.synchronize()
        
        compute_time = time.time() - compute_start
        total_time = time.time() - start_time
        
        print(f"✅ 矩阵乘法耗时: {compute_time:.3f}秒")
        print(f"📊 总耗时: {total_time:.3f}秒")
        
        # 性能评级（基于 RTX 3080 预期性能）
        if compute_time < 0.5:
            print("🏆 性能评级: 优秀 (GPU加速正常)")
        elif compute_time < 1.0:
            print("👍 性能评级: 良好")
        else:
            print("⚠️  性能评级: 一般 (可能存在问题)")
            
        # 清理显存
        del a, b, c
        torch.cuda.empty_cache()
        
    else:
        print("\n❌ 未检测到CUDA支持")
        print("💡 解决建议:")
        print("1. 确保安装了 NVIDIA 驱动")
        print("2. 安装 CUDA Toolkit")
        print("3. 安装支持CUDA的PyTorch版本")

def test_whisper_requirements():
    """测试 Whisper 相关依赖"""
    print(f"\n🎵 Whisper 依赖检测")
    print("=" * 50)
    
    try:
        import whisper
        print("✅ OpenAI Whisper: 已安装")
        
        # 获取可用模型
        models = whisper.available_models()
        print(f"📦 可用模型: {', '.join(models)}")
        
        # 推荐配置
        print(f"\n💡 RTX 3080 推荐配置 (已升级):")
        print(f"   🚀 当前使用: medium 模型")
        print(f"   💾 medium模型显存需求: ~5GB")
        print(f"   💾 RTX 3080 显存总量: 10GB (绰绰有余)")
        print(f"   ⚡ 预期速度: 比 small 模型慢 20-30%，但质量显著提升")
        print(f"   🎯 质量提升: 更准确的中文识别和标点符号")
        print(f"   📊 适合场景: 专业转录、重要会议记录")
        
    except ImportError:
        print("❌ OpenAI Whisper: 未安装")
        print("💡 请运行: pip install openai-whisper")

if __name__ == "__main__":
    print("🚀 RTX 3080 GPU 性能检测工具")
    print("时间:", time.strftime("%Y-%m-%d %H:%M:%S"))
    print()
    
    check_gpu_status()
    test_whisper_requirements()
    
    print(f"\n✨ 检测完成!")
    print("=" * 50)